//==========================================================================
// Copyright (c) 2000-2008,  Elastos, Inc.  All Rights Reserved.
//==========================================================================
#include <aura.h>
#include <string.h>
#include <stdlib.h>
#include "win32module.h"
#include "ntimage.h"
#include "cgmodule.h"
#include "pub.h"

extern DLinkNode s_ModuleList;
extern Aura_Handle g_hAttachLock;
extern Aura_Handle s_hModuleLock;

Aura_ECode GlobalLoadModuleW(
        const Aura_PWChar wszName,
        Aura_UInt32 dwFlags,
        ModuleFileType  modFileType,
        GModule ** ppModule);

GModule * FindModuleByImageBase(Aura_UInt32 uImageBase);

Aura_UInt32 GetExportAddressByName(
        GModule * pModule,
        const CHAR * pszName);
        
Aura_UInt32 GetExportAddressByOrdinal(
        GModule * pModule,
        Aura_UInt32 uOrdinal);        
        
       
AURA_API AuraModule_Load(
        /* [in] */ Aura_PWChar name,
        /* [in] */ Aura_ModuleLoadFlags flags,
        /* [out] */ Aura_Handle * phModule)
{
    GModule * pModule;
    Aura_ECode ec = AURA_NOERROR;
    
    Aura_PWChar wszLoadModuleName;
    wszLoadModuleName = name;
    
    ModuleFileType modFileType = ModuleFileType_DLL;
    
    if(!wcsstr((wchar_t *)name, L".dll") && wcsstr((wchar_t *)name, L".exe")) {
        modFileType = ModuleFileType_EXE;
    }
    
    ec = GlobalLoadModuleW(wszLoadModuleName, (Aura_UInt32)flags, modFileType, &pModule); 
    if(AURA_FAILED(ec)){
        if (ec == AURA_E_FILE_NOT_FOUND && modFileType == ModuleFileType_DLL) {
            return Win32Module_Load(name, flags, phModule);
        }
        return ec;
    }
    
    Lock(g_hAttachLock);
    if(!ProcessAttach(pModule, AURA_NULL)){
        Aura_printf("Error: ProcessAttach failed! file:%s. line:%d. \n", __FILE__, __LINE__);
        Unlock(g_hAttachLock);
        return AURA_E_INVALID_OPERATION;
    }
    Unlock(g_hAttachLock);
    
    *phModule = (Aura_Handle)pModule->m_uImageBase;
    return ec;
}


AURA_API AuraModule_Unload(
        /* [in] */ Aura_Handle hModule)
{   
    GModule * pModule;
    pModule = FindModuleByImageBase((Aura_UInt32)hModule);
    if(AURA_NULL == pModule){
        return Win32Module_Unload(hModule);
    }
    pModule->Release(); //Release the FindModuleByImageBase func
    pModule->Release();
    return AURA_NOERROR;
}


AURA_API AuraModule_Close(
        /* [in] */ Aura_Handle hModule)
{   
    GModule * pModule;
    pModule = FindModuleByImageBase((Aura_UInt32)hModule);
    if(AURA_NULL == pModule){
        return  Win32Module_Close(hModule);
    }
    
    pModule->Release();
    return AURA_NOERROR;
}


AURA_API AuraModule_GetFullFileName(
        /* [in] */ Aura_Handle hModule,
        /* [out] */Aura_PWChar pFullNameBuf,
        /* [in] */ Aura_Int32 bufSize)
{   
    GModule * pModule;
   
    if((AURA_NULL == hModule)||(AURA_NULL == pFullNameBuf)){
        return AURA_E_INVALID_ARGUMENT;
    }
    
    pModule = FindModuleByImageBase((Aura_UInt32)hModule);
    if(AURA_NULL == pModule){
        return Win32Module_GetFullFileName(hModule, pFullNameBuf, bufSize);
    }
    
    if((Aura_UInt32)bufSize < wcslen((const wchar_t*)pModule->m_wszFullPath)){
        return AURA_E_INVALID_ARGUMENT;
    }
    
    wcscpy((wchar_t*)pFullNameBuf, (const wchar_t*)pModule->m_wszFullPath);
    pModule->Release();
    
    return AURA_NOERROR;
}


AURA_API AuraModule_GetProcAddressByName(
        /* [in] */ Aura_Handle hModule,
        /* [in] */ Aura_PChar pszName,
        /* [out] */ Aura_Address * pProcAddress)
{   
    GModule * pModule;
    Aura_UInt32 entry;

    if(pszName == AURA_NULL ||(AURA_NULL == pProcAddress)){
        return AURA_E_INVALID_ARGUMENT;
    }

    pModule = FindModuleByImageBase((Aura_UInt32)hModule);
    if(AURA_NULL == pModule){
        return Win32Module_GetProcAddressByName(hModule, pszName, pProcAddress);
    }

    entry = GetExportAddressByName(pModule, (char *)pszName);
    if(0 == entry){
        return AURA_E_DOES_NOT_EXIST;
    }
    
    *pProcAddress = (Aura_Address)entry;

    pModule->Release();

    return AURA_NOERROR;
}


AURA_API AuraModule_GetProcAddressByOrdinal(
        /* [in] */ Aura_Handle hModule,
        /* [in] */ Aura_Int32 ordinal,
        /* [out] */ Aura_Address * pProcAddress)       
{      
    GModule * pModule;
    Aura_UInt32 entry;

    if(AURA_NULL == pProcAddress){
        return AURA_E_INVALID_ARGUMENT;
    }
    
    pModule = FindModuleByImageBase((Aura_UInt32)hModule);
    if(AURA_NULL == pModule){
        return Win32Module_GetProcAddressByOrdinal(hModule, ordinal, pProcAddress);
    }

    entry = GetExportAddressByOrdinal(pModule, ordinal);
    if(0 == entry){
        return AURA_E_DOES_NOT_EXIST;
    }
    
    *pProcAddress = (Aura_Address)entry;

    pModule->Release();

    return AURA_NOERROR;
}


AURA_API AuraModule_GetResourceSection(
        /* [in] */ Aura_Address imageBaseAddr,
        /* [out] */ Aura_Address * pSectBaseAddr)
{      
    GModule * pModule;

    if(AURA_NULL == pSectBaseAddr){
        return AURA_E_INVALID_ARGUMENT;
    }

    pModule = FindModuleByImageBase((Aura_UInt32)imageBaseAddr);
    if(AURA_NULL == pModule){
        return Win32Module_GetResourceSection(imageBaseAddr, pSectBaseAddr);
    }
    
    *pSectBaseAddr = imageBaseAddr + (Aura_Address)pModule->m_pResource->m_uVAddr;
       
    pModule->Release();

    return AURA_NOERROR;
}


struct ModuleSnapshot : public DLinkNode {  
    Aura_UInt32          m_uImageBase;
    ModuleSnapshot(Aura_UInt32 uImageBase) : m_uImageBase(uImageBase) { }
};

struct ModuleSnapshotList{
    DLinkNode      m_SnapshotList;
    DLinkNode     *m_pCurrent;
    ModuleSnapshotList()
    {
        m_SnapshotList.Initialize();
        m_pCurrent = &m_SnapshotList;
    }
};


Aura_ECode CreateModuleSnapShotList(DLinkNode *pSnapshotList, DLinkNode *pModuleList)
{
    ModuleSnapshot *pModuleSnapshot = AURA_NULL;
    GModule *pModule;

    Lock(s_hModuleLock);
    ForEachDLinkNode(GModule *, pModule, pModuleList) {
        pModuleSnapshot = new ModuleSnapshot(pModule->m_uImageBase);
        if(!pModuleSnapshot){
            Unlock(s_hModuleLock);
            return AURA_E_OUT_OF_MEMORY;
        }
        pSnapshotList->InsertNext(pModuleSnapshot);
    }

    Unlock(s_hModuleLock);
    return AURA_NOERROR;
}


Aura_ECode DestroyModuleSnapShotList(DLinkNode *pSnapshotList)
{
    ModuleSnapshot *pModuleSnapshot = AURA_NULL;

    Lock(s_hModuleLock);
    while(pSnapshotList != pSnapshotList->Next()){
        pModuleSnapshot = (ModuleSnapshot *)pSnapshotList->Next();
        pModuleSnapshot->Detach();
        delete pModuleSnapshot;
    }
    Unlock(s_hModuleLock);
    return AURA_NOERROR;
}


AURA_API AuraModule_CreateModuleSnapshot(
        /* [in] */ Aura_ProcessId processId,
        /* [out] */ Aura_Handle * phSnapshot)
{    
    return Win32Module_CreateModuleSnapshot(processId, phSnapshot);
   
    Aura_ProcessId currentPId;
    ModuleSnapshotList *pSnapshotList = AURA_NULL;

    Aura_ECode ec;
    if(AURA_NULL == phSnapshot){
        return AURA_E_INVALID_ARGUMENT;
    }

    ec = AuraProcess_GetCurrentId(&currentPId);
    if(AURA_FAILED(ec)){
        return ec;
    }

    if(currentPId != processId){
        return AURA_E_NOT_IMPLEMENTED;
    }

    pSnapshotList = new ModuleSnapshotList();
    if (!pSnapshotList) {
        return AURA_E_OUT_OF_MEMORY;
    }

    ec = CreateModuleSnapShotList(&pSnapshotList->m_SnapshotList, &s_ModuleList);
    if(AURA_FAILED(ec)){
        return ec;
    }

    *phSnapshot = (Aura_Handle)pSnapshotList;

    return AURA_NOERROR;
}


AURA_API AuraModule_CloseModuleSnapshot(
        /* [in] */ Aura_Handle hSnapshot)
{
    return Win32Module_CloseModuleSnapshot(hSnapshot);
    
    if(AURA_NULL == hSnapshot){
        return AURA_E_INVALID_ARGUMENT;
    }

    return DestroyModuleSnapShotList(&((ModuleSnapshotList *)hSnapshot)->m_SnapshotList);
}


AURA_API AuraModule_GetModuleNode(
        /* [in] */ Aura_Handle hSnapshot,
        /* [in] */ Aura_Bool bFirstNode,
        /* [out] */ Aura_Bool * pbExistedMore,
        /* [out] */ Aura_Handle * phModule)
{
    return Win32Module_GetModuleNode( 
        hSnapshot, bFirstNode, pbExistedMore, phModule);
    
    ModuleSnapshotList *pSnapshotList = AURA_NULL;
    if((AURA_NULL == hSnapshot)||(AURA_NULL == pbExistedMore)||(AURA_NULL == phModule)){
        return AURA_E_INVALID_ARGUMENT;
    }

    pSnapshotList = (ModuleSnapshotList *)hSnapshot;

    if(bFirstNode){
        pSnapshotList->m_pCurrent = pSnapshotList->m_SnapshotList.Next();
    } else {
        pSnapshotList->m_pCurrent = pSnapshotList->m_pCurrent->Next();
    }

    if(pSnapshotList->m_pCurrent == &pSnapshotList->m_SnapshotList){
        *pbExistedMore = AURA_FALSE;
        *phModule = AURA_NULL;
    } else {
        *pbExistedMore = AURA_TRUE;
        *phModule = (Aura_Handle)((ModuleSnapshot *)pSnapshotList->m_pCurrent)->m_uImageBase;
    }

    return AURA_NOERROR;
}


AURA_API AuraModule_GetInformation(
        /* [in] */ Aura_ProcessId processId,
        /* [in] */ Aura_Handle hModule,
        /* [out] */ Aura_Address * pBaseAddr,
        /* [out] */ Aura_Address * pEntryAddr,
        /* [out] */ Aura_MemorySize * pSizeOfImage)
{
    Aura_ProcessId CurrentProcessId;
    GModule * pModule;

    AuraProcess_GetCurrentId(&CurrentProcessId);

    if(CurrentProcessId != processId){
        return AURA_E_NOT_IMPLEMENTED;
    }

    pModule = FindModuleByImageBase((Aura_UInt32)hModule);
    if(AURA_NULL == pModule){
        return Win32Module_GetInformation(processId, hModule, \
                pBaseAddr, pEntryAddr, pSizeOfImage); 
    }

    if(AURA_NULL != pBaseAddr){
        *pBaseAddr = (Aura_Address)pModule->m_uImageBase;
    }

    if(AURA_NULL != pEntryAddr){
        *pEntryAddr = (Aura_Address)pModule->m_uEntryPoint;
    }

    if(AURA_NULL != pSizeOfImage){
        *pSizeOfImage = (Aura_MemorySize)pModule->m_uTotalSize;
    }

    pModule->Release();

    return AURA_NOERROR;
}


PIMAGE_RESOURCE_DIRECTORY FindResDir(PIMAGE_RESOURCE_DIRECTORY prd,
                                     DWORD dwResBase,
                                     Aura_PWChar pName)
{
    PIMAGE_RESOURCE_DIRECTORY_ENTRY prde;
    prde = (PIMAGE_RESOURCE_DIRECTORY_ENTRY)((DWORD)prd +
                sizeof(IMAGE_RESOURCE_DIRECTORY));

    if (*pName == '#') {
        int id = _wtoi((const wchar_t*)(pName + 1));
        prde += prd->NumberOfNamedEntries;
        for (int i = 0; i < prd->NumberOfIdEntries; i++) {
            if (prde->Id == id) {
                if (prde->DataIsDirectory) {
                    return (PIMAGE_RESOURCE_DIRECTORY)
                             ((char*)dwResBase + prde->OffsetToDirectory);
                }
            }
            prde++;
        }
    }
    else {
        PIMAGE_RESOURCE_DIRECTORY_STRING prds = AURA_NULL;
        /* loop through all resource directory entry types */
        for (int i = 0; i < prd->NumberOfNamedEntries; i++) {
            if (prde->NameIsString) {
                prds = (PIMAGE_RESOURCE_DIRECTORY_STRING)
                                ((char*)dwResBase + prde->NameOffset);

                if (_wcsnicmp((const wchar_t*)(prds->NameString),  \
                        (const wchar_t*)pName, prds->Length) == 0) {
                    if (prde->DataIsDirectory) {
                        return (PIMAGE_RESOURCE_DIRECTORY)
                                 ((char*)dwResBase + prde->OffsetToDirectory);
                    }
                }
            }
            prde++;
        }
    }

    return AURA_NULL;
}


PIMAGE_RESOURCE_DATA_ENTRY FindResEntry(PIMAGE_RESOURCE_DIRECTORY prd,
                                        DWORD dwResBase)
{
    PIMAGE_RESOURCE_DIRECTORY_ENTRY prde;
    prde = (PIMAGE_RESOURCE_DIRECTORY_ENTRY)((DWORD)prd +
                sizeof(IMAGE_RESOURCE_DIRECTORY));

    if (!prde->DataIsDirectory) {
        return (PIMAGE_RESOURCE_DATA_ENTRY)
                        ((char*)dwResBase + prde->OffsetToData);
    }
    else {
        return AURA_NULL;
    }
}


AURA_API AuraModule_GetResource(
    /* [in] */ Aura_Handle hModule,
    /* [in] */ Aura_PWChar wsName,
    /* [in] */ Aura_PWChar wsType,
    /* [out] */ Aura_Address * pAddress,
    /* [out] */ Aura_MemorySize * pSize)
{
    GModule * pModule;
    PIMAGE_RESOURCE_DIRECTORY  prd, prd2, prd3;
    PIMAGE_RESOURCE_DATA_ENTRY pdata;
    DWORD dwBase, dwResBase;

    Aura_ECode ec = AURA_NOERROR;

    if (AURA_NULL == pSize || AURA_NULL == pAddress
        || wsName == AURA_NULL || wsType == AURA_NULL) {
        return AURA_E_INVALID_ARGUMENT;
    }

    pModule = FindModuleByImageBase((Aura_UInt32)hModule);
    if(AURA_NULL == pModule){
        return Win32Module_GetResource(hModule, wsName, wsType, pAddress, pSize);
    }

    dwBase = pModule->m_uImageBase;

    ec = AuraModule_GetResourceSection((Aura_Address)dwBase, (Aura_Address *)&dwResBase);
    if (AURA_FAILED(ec)) {
        return ec;
    }

    prd = (PIMAGE_RESOURCE_DIRECTORY)dwResBase;
    /* set pointer to first resource type entry */
    prd2 = FindResDir(prd, dwResBase, (Aura_PWChar)wsType);
    if (prd2 == AURA_NULL) return AURA_E_DOES_NOT_EXIST;

    prd3 = FindResDir(prd2, dwResBase, (Aura_PWChar)wsName);
    if (prd3 == AURA_NULL) return AURA_E_DOES_NOT_EXIST;

    pdata = FindResEntry(prd3, dwResBase);
    if (pdata == AURA_NULL) return AURA_E_DOES_NOT_EXIST;

    *pAddress = (Aura_Address)((char*)dwBase + pdata->OffsetToData);
    *pSize = pdata->Size;

    pModule->Release();

    return AURA_NOERROR;
}
